package com.lzx.RandomForest;


import java.io.IOException;
import java.util.*;

/**
 * @author: 如寄
 * @version: v1.0
 * @date:2021/5/11
 */
public class CreateTree {
    private StringBuilder tree;//决策树
    private Map<Integer, Boolean> visit = new HashMap<>();//记录已经访问过的特征
    private List<String> features;

    /**
     * 生成决策树
     *
     * @param testData:测试数据集
     * @throws IOException
     * @throws InterruptedException
     * @return: 以前序遍历的一棵决策树模型
     */
    public String buildTree(Map<String, List<Map<Integer, Boolean>>> trainData, Map<String, List<Map<Integer, Boolean>>> testData, List<String> featherList) throws IOException, InterruptedException {
        features = featherList;
        TreeNode node = dfs(trainData);
        cut(node, testData);
        tree = new StringBuilder();
        pre(node, tree);
        System.out.println(tree);
        System.out.println("$$$");
        return tree.toString();
    }

    /**
     * 剪枝、防止过拟合
     *
     * @param node     :决策树
     * @param testData :训练数据
     * @return :节点剪枝前和剪枝后的正确分类的最大值
     */
    private int cut(TreeNode node, Map<String, List<Map<Integer, Boolean>>> testData) {
        if (node.left == null && node.right == null) {
            return testData.containsKey(node.getForecast()) ? testData.get(node.getForecast()).size() : 0;
        }
        //计算剪枝后的正确数
        int a = testData.containsKey(node.getForecast()) ? testData.get(node.getForecast()).size() : 0;
        //计算不剪枝的正确数
        List<Map<String, List<Map<Integer, Boolean>>>> divideData = divideData(node.index, testData);
        int l = cut(node.left, divideData.get(0));
        int r = cut(node.right, divideData.get(1));
        if (a >= (l + r)) {//判断是否剪枝
            node.node = node.getForecast();
            node.right = null;
            node.left = null;
        }
        return Math.max(a, l + r);
    }


    /**
     * 前序遍历决策树
     *
     * @param node:决策树
     */
    private void pre(TreeNode node, StringBuilder tree) {
        if (node.right == null && node.left == null) {
            tree.append(node.node + " # # ");
            return;
        }
        tree.append(features.get(node.index) + " ");
        pre(node.left, tree);
        pre(node.right, tree);
    }

    /**
     * 前序遍历构造决策树
     *
     * @param data:训练数据
     * @return :决策树根节点
     * @throws IOException
     * @throws InterruptedException
     */
    private TreeNode dfs(Map<String, List<Map<Integer, Boolean>>> data) throws
            IOException, InterruptedException {
        String firstType = (String) ((data.keySet().toArray())[0]);
        List<Map<Integer, Integer>> featureList = new ArrayList<>(data.get(firstType).size() * data.size() * 32);//每个类型的特征集
        Map<Integer, Integer> wholeMap = new HashMap<>();//数据集中所有特征的集
        List<Integer> eachTypeNumber = new ArrayList<>();//每个类型数据的数量
        int allSum = 0;//所有数据的数量
        if (data.size() == 1) { //数据类型只有一个的时候树停止生长，以此类型作为叶子节点
            String type = getType(data);
            TreeNode node = new TreeNode(type);
            node.setForecast(type);
            return node;
        }
        for (Map.Entry<String, List<Map<Integer, Boolean>>> entry1 : data.entrySet()) {
            List<Map<Integer, Boolean>> list = entry1.getValue();
            Map<Integer, Integer> res = new HashMap<Integer, Integer>(list.size() * 32);
            for (Map<Integer, Boolean> map : list) {
                for (Map.Entry<Integer, Boolean> entry : map.entrySet()) {
                    res.put(entry.getKey(), res.getOrDefault(entry.getKey(), 0) + 1);
                    wholeMap.put(entry.getKey(), wholeMap.getOrDefault(entry.getKey(), 0) + 1);
                }
            }
            featureList.add(res);
            allSum += list.size();
            eachTypeNumber.add(list.size());
        }
        Integer feature = null;
        Map<Integer, Double> temp = getGain(wholeMap, eachTypeNumber, featureList, allSum);
        feature = getRandomFeature(temp);
        String type = getType(data);
        if (feature == null) {
            TreeNode node = new TreeNode(type);
            node.setForecast(type);
            return node;
        }
        TreeNode node = new TreeNode(feature);
        node.setForecast(type);
        visit.put(feature, true);
        List<Map<String, List<Map<Integer, Boolean>>>> divideData = divideData(feature, data);
        data = null;
        wholeMap = null;
        featureList = null;
        Map<String, List<Map<Integer, Boolean>>> left = divideData.get(0);
        Map<String, List<Map<Integer, Boolean>>> right = divideData.get(1);
        node.left = dfs(left);
        if (right.size() == 0) {
            node.right = new TreeNode(type);
            node.right.setForecast(type);
        } else node.right = dfs(right);
        visit.remove(feature);
        return node;
    }


    /**
     * 获取随机特征
     *
     * @param temp :特征集
     * @return 随机选出的特征
     */
    private Integer getRandomFeature(Map<Integer, Double> temp) {
        List<Map.Entry<Integer, Double>> lst = new ArrayList<Map.Entry<Integer, Double>>(temp.entrySet());
        Random random = new Random();
        int i = (int) StrictMath.pow(lst.size(), 1.0/3.0);
        if (i <= 0) {
            return null;
        }
        i = random.nextInt(i) + 1;
        return findKthBig(lst, i);
    }


    /**
     * 计算信息增益
     *
     * @param wholeMap                :所有 特征-特征数量 集
     * @param eachTypeNumber:每个类型的新闻数
     * @param featureList:每个类型的       特征-特征数量 集
     * @param allSum:|新闻数量集|
     * @return:每个特征-信息增益集
     */
    private Map<Integer, Double> getGain
    (Map<Integer, Integer> wholeMap, List<Integer> eachTypeNumber, List<Map<Integer, Integer>> featureList,
     int allSum) {
        Map<Integer, Double> temp = new HashMap<>(allSum * 32);
        double ED = getEnt(allSum, eachTypeNumber);
        for (Map.Entry<Integer, Integer> entry : wholeMap.entrySet()) {
            double e1 = getLeEnt(entry.getValue(), featureList, entry.getKey());
            double e2 = getReEnt(allSum - entry.getValue(), featureList, entry.getKey(), eachTypeNumber);
            //计算信息增益gain
            double gain = ED - ((entry.getValue() * 1.0 / allSum) * e1 + ((allSum - entry.getValue()) * 1.0 / allSum) * e2);
            if (!visit.containsKey(entry.getKey()) && !entry.getKey().equals("\n") && !entry.getKey().equals("")) {
                temp.put(entry.getKey(), gain);
            }
        }
        return temp;
    }


    /**
     * 获得次节点上的数据数量最多的类型
     *
     * @param ans:节点的数据集
     * @return：节点上的数据数量最多的类型
     */
    private String getType(Map<String, List<Map<Integer, Boolean>>> ans) {
        String type = "其它";
        int temp = 0;
        for (Map.Entry<String, List<Map<Integer, Boolean>>> entry : ans.entrySet()) {
            if (entry.getValue().size() > temp) {
                type = entry.getKey();
                temp = entry.getValue().size();
            }
        }
        return type;
    }


    /**
     * 获得此节点上包含特征feature的数据集
     *
     * @param feature ：特征
     * @param res     ：节点的数据集
     * @return：包含特征feature的数据集
     */
    private List<Map<String, List<Map<Integer, Boolean>>>> divideData(Integer
                                                                              feature, Map<String, List<Map<Integer, Boolean>>> res) {
        List<Map<String, List<Map<Integer, Boolean>>>> list = new ArrayList<>();
        Map<String, List<Map<Integer, Boolean>>> left = new HashMap<>();
        Map<String, List<Map<Integer, Boolean>>> right = new HashMap<>();
        for (Map.Entry<String, List<Map<Integer, Boolean>>> entry : res.entrySet()) {
            List<Map<Integer, Boolean>> mapList = entry.getValue();
            List<Map<Integer, Boolean>> leftList = new ArrayList<>();
            List<Map<Integer, Boolean>> rightList = new ArrayList<>();
            for (Map<Integer, Boolean> map : mapList) {
                if (map.containsKey(feature)) {
                    leftList.add(map);
                } else rightList.add(map);
            }
            if (leftList.size() > 0) {
                left.put(entry.getKey(), leftList);
            }
            if (rightList.size() > 0) {
                right.put(entry.getKey(), rightList);
            }
        }
        list.add(left);
        list.add(right);
        return list;
    }


    /**
     * 计算不包含特征feature的信息熵
     *
     * @param value:不包含特征feature的数据数量
     * @param listS：每个类型的特征集
     * @param feature：特征
     * @param allList：每个类型的数量
     * @return:不包含特征feature的信息熵
     */
    private static double getReEnt(int value, List<Map<Integer, Integer>> listS, Integer
            feature, List<Integer> allList) {
        if (value == 0)
            return 0;
        double res = 0;
        int i = 0;
        for (Map<Integer, Integer> map : listS) {
            if (map.containsKey(feature)) {
                res += (allList.get(i) - map.get(feature) + 1) * 1.0 / value * Math.log((allList.get(i) - map.get(feature) + 1) * 1.0 / value);
            } else res += allList.get(i) * 1.0 / value * Math.log(allList.get(i) * 1.0 / value);
            i++;
        }
        return -res;
    }


    /**
     * 计算包含特征feature的信息熵
     *
     * @param value：包含特征feature的数据数量
     * @param listS：每个类型的特征集
     * @param feature：特征
     * @return：包含特征feature的信息熵
     */
    private static double getLeEnt(int value, List<Map<Integer, Integer>> listS, Integer feature) {
        double res = 0;
        for (Map<Integer, Integer> map : listS) {
            if (map.containsKey(feature)) {
                res += map.get(feature) * 1.0 / value * Math.log(map.get(feature) * 1.0 / value);
            }
        }
        return -res;
    }


    /**
     * 计算数据集的信息熵
     *
     * @param allSum:|数据集|
     * @param allList：每个类型的数量
     * @return：数据的信息熵
     */
    private static double getEnt(int allSum, List<Integer> allList) {
        double res = 0;
        for (Integer integer : allList) {
            res += integer * 1.0 / allSum * Math.log(integer * 1.0 / allSum);
        }
        return -res;
    }

    private Integer findKthBig(List<Map.Entry<Integer, Double>> list, int k) {
        return findK(list, k, 0, list.size() - 1);
    }

    private Integer findK(List<Map.Entry<Integer, Double>> list, int k, int l, int r) {
        if (l > r)
            return null;
        int left = l, right = r;
        Map.Entry<Integer, Double> temp = list.get(l);
        while (left < right) {
            while (left < right && list.get(right).getValue() <= temp.getValue()) {
                right--;
            }
            list.set(left, list.get(right));
            while (left < right && list.get(left).getValue() >= temp.getValue()) {
                left++;
            }
            list.set(right, list.get(left));
        }
        list.set(right, temp);
        if (right == k - 1)
            return temp.getKey();
        else if (right > k - 1)
            return findK(list, k, l, right - 1);
        else return findK(list, k, right + 1, r);
    }


}
